{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 09: Stokes-Darcy equations\n",
    "\n",
    "Coupled mixed formulation using Lagrange multiplier from\n",
    "```\n",
    "Layton, Schieweck, Yotov, Coupling fluid flow with porous media flow. SINUM 2003, DOI:10.1137/S0036142901392766\n",
    "```\n",
    "\n",
    "defined on the interface\n",
    "\n",
    "\\begin{equation*}\n",
    "\\begin{aligned}\n",
    "u_S: \\text{Stokes velocity in }H^1(\\Omega_S)\\\\\n",
    "u_D: \\text{Darcy velocity in }H(\\text{div},\\Omega_D)\\\\\n",
    "p_S: \\text{Stokes pressure in }L^2(\\Omega_S)\\\\\n",
    "p_D: \\text{Darcy pressure in }L^2(\\Omega_D)\\\\\n",
    "\\end{aligned}\n",
    "\\end{equation*}\n",
    "\n",
    "transmission conditions:\n",
    "\n",
    "\\begin{equation*}\n",
    "\\begin{cases}\n",
    "u_S \\cdot n_S + u_D \\cdot n_D = 0\\\\\n",
    "-(2 \\mu \\varepsilon(u_S) - p_S I) n_S \\cdot n_S = pD\\\\\n",
    "-(2 \\mu \\varepsilon(u_S) - p_S I) t_S \\cdot n_S = \\alpha \\mu k^{-0.5} u_S \\cdot t_S\n",
    "\\end{cases}\n",
    "\\end{equation*}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from mpi4py import MPI\n",
    "from petsc4py import PETSc\n",
    "from ufl import as_vector, avg, div, dot, FacetNormal, grad, inner, Measure, sym, TestFunction, TrialFunction\n",
    "from dolfinx import (Constant, DirichletBC, Function, FunctionSpace, MeshTags, RectangleMesh,\n",
    "                     VectorFunctionSpace)\n",
    "from dolfinx.cpp.mesh import GhostMode\n",
    "from dolfinx.fem import locate_dofs_topological\n",
    "from dolfinx.mesh import locate_entities\n",
    "from multiphenicsx.fem import (assemble_matrix_block, assemble_scalar, assemble_vector_block,\n",
    "                               BlockVecSubVectorWrapper, create_vector_block, DofMapRestriction)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct mesh\n",
    "if MPI.COMM_WORLD.size > 1:\n",
    "    mesh_ghost_mode = GhostMode.shared_facet  # shared_facet ghost mode is required by dS\n",
    "else:\n",
    "    mesh_ghost_mode = GhostMode.none\n",
    "mesh = RectangleMesh(MPI.COMM_WORLD, [np.array([-1.0, -2.0, 0.0]), np.array([1.0, 2.0, 0.0])],\n",
    "                     [50, 100], ghost_mode=mesh_ghost_mode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper functions for boundaries and subdomains marking\n",
    "eps = np.finfo(float).eps\n",
    "\n",
    "\n",
    "def near(x, a):\n",
    "    return abs(x - a) < eps\n",
    "\n",
    "\n",
    "def above(x, a):\n",
    "    return x > a - eps\n",
    "\n",
    "\n",
    "def below(x, a):\n",
    "    return x < a + eps\n",
    "\n",
    "\n",
    "def between(x, interval):\n",
    "    return np.logical_and(above(x, interval[0]), below(x, interval[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set subdomains\n",
    "def darcy(x):\n",
    "    return below(x[1], 0.0)\n",
    "\n",
    "\n",
    "def stokes(x):\n",
    "    return above(x[1], 0.0)\n",
    "\n",
    "\n",
    "darcy.id = 10\n",
    "stokes.id = 13\n",
    "\n",
    "subdomains_entities = dict()\n",
    "subdomains_values = dict()\n",
    "for subdomain in (darcy, stokes):\n",
    "    subdomains_entities[subdomain.id] = locate_entities(\n",
    "        mesh, mesh.topology.dim, subdomain)\n",
    "    subdomains_values[subdomain.id] = np.full(\n",
    "        subdomains_entities[subdomain.id].shape, subdomain.id, dtype=np.intc)\n",
    "subdomains = MeshTags(mesh, mesh.topology.dim,\n",
    "                      np.hstack(list(subdomains_entities.values())),\n",
    "                      np.hstack(list(subdomains_values.values())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set boundaries and interface\n",
    "def outlet(x):\n",
    "    return near(x[1], -2.0)\n",
    "\n",
    "\n",
    "def inlet(x):\n",
    "    return near(x[1], 2.0)\n",
    "\n",
    "\n",
    "def interface(x):\n",
    "    return near(x[1], 0.0)\n",
    "\n",
    "\n",
    "def wall_stokes(x):\n",
    "    return np.logical_or(\n",
    "        np.logical_and(near(x[0], 1.0), between(x[1], (0.0, 2.0))),\n",
    "        np.logical_and(near(x[0], -1.0), between(x[1], (0.0, 2.0)))\n",
    "    )\n",
    "\n",
    "\n",
    "def wall_darcy(x):\n",
    "    return np.logical_or(\n",
    "        np.logical_and(near(x[0], 1.0), between(x[1], (-2.0, 0.0))),\n",
    "        np.logical_and(near(x[0], -1.0), between(x[1], (-2.0, 0.0)))\n",
    "    )\n",
    "\n",
    "\n",
    "outlet.id = 14\n",
    "inlet.id = 15\n",
    "interface.id = 16\n",
    "wall_stokes.id = 17\n",
    "wall_darcy.id = 18\n",
    "\n",
    "boundaries_entities = dict()\n",
    "boundaries_values = dict()\n",
    "for boundary in (outlet, inlet, interface, wall_stokes, wall_darcy):\n",
    "    boundaries_entities[boundary.id] = locate_entities(\n",
    "        mesh, mesh.topology.dim - 1, boundary)\n",
    "    boundaries_values[boundary.id] = np.full(\n",
    "        boundaries_entities[boundary.id].shape, boundary.id, dtype=np.intc)\n",
    "boundaries = MeshTags(mesh, mesh.topology.dim - 1,\n",
    "                      np.hstack(list(boundaries_entities.values())),\n",
    "                      np.hstack(list(boundaries_values.values())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define associated measures\n",
    "dx = Measure(\"dx\")(subdomain_data=subdomains)\n",
    "ds = Measure(\"ds\")(subdomain_data=boundaries)\n",
    "dS = Measure(\"dS\")(subdomain_data=boundaries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify the domain size\n",
    "one = Constant(mesh, 1.0)\n",
    "areaD = mesh.mpi_comm().allreduce(assemble_scalar(one * dx(darcy.id)), op=MPI.SUM)\n",
    "areaS = mesh.mpi_comm().allreduce(assemble_scalar(one * dx(stokes.id)), op=MPI.SUM)\n",
    "lengthI = mesh.mpi_comm().allreduce(assemble_scalar(one * dS(interface.id)), op=MPI.SUM)\n",
    "print(\"area(Omega_D) = \", areaD)\n",
    "print(\"area(Omega_S) = \", areaS)\n",
    "print(\"length(Sigma) = \", lengthI)\n",
    "assert np.isclose(areaD, 4.)\n",
    "assert np.isclose(areaS, 4.)\n",
    "assert np.isclose(lengthI, 2.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Normal and tangent\n",
    "n = FacetNormal(mesh)\n",
    "t = as_vector([n[1], -n[0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function spaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P2v = VectorFunctionSpace(mesh, (\"CG\", 2))  # uS\n",
    "BDM1 = FunctionSpace(mesh, (\"BDM\", 1))  # uD. It can also be RT\n",
    "P1 = FunctionSpace(mesh, (\"CG\", 1))  # pS\n",
    "P0 = FunctionSpace(mesh, (\"DG\", 0))  # pD\n",
    "Pt = FunctionSpace(mesh, (\"DGT\", 1))  # la"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Restrictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define restrictions\n",
    "dofs_P2v = locate_dofs_topological(P2v, subdomains.dim, subdomains_entities[stokes.id])\n",
    "dofs_BDM1 = locate_dofs_topological(BDM1, subdomains.dim, subdomains_entities[darcy.id])\n",
    "dofs_P1 = locate_dofs_topological(P1, subdomains.dim, subdomains_entities[stokes.id])\n",
    "dofs_P0 = locate_dofs_topological(P0, subdomains.dim, subdomains_entities[darcy.id])\n",
    "dofs_Pt = locate_dofs_topological(Pt, boundaries.dim, boundaries_entities[interface.id])\n",
    "restriction = [DofMapRestriction(V.dofmap, dofs) for (V, dofs) in zip(\n",
    "               (P2v, BDM1, P1, P0, Pt), (dofs_P2v, dofs_BDM1, dofs_P1, dofs_P0, dofs_Pt))]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Trial and test functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define trial and test functions\n",
    "uS, uD, pS, pD, la = [TrialFunction(V) for V in (P2v, BDM1, P1, P0, Pt)]\n",
    "vS, vD, qS, qD, xi = [TestFunction(V) for V in (P2v, BDM1, P1, P0, Pt)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Problem data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mu = 1.\n",
    "alpha = 1.\n",
    "k = 1.\n",
    "fS = Constant(mesh, (0., 0.))\n",
    "fD = fS\n",
    "gS = Constant(mesh, 0.)\n",
    "gD = gS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Weak formulation and boundary conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def epsilon(vec):\n",
    "    return sym(grad(vec))\n",
    "\n",
    "\n",
    "AS = (2.0 * mu * inner(epsilon(uS), epsilon(vS)) * dx(stokes.id)\n",
    "      + mu * alpha * pow(k, -0.5) * dot(uS(\"+\"), t(\"+\")) * dot(vS(\"+\"), t(\"+\")) * dS(interface.id))\n",
    "\n",
    "AD = mu / k * dot(uD, vD) * dx(darcy.id)\n",
    "\n",
    "B1St = - pS * div(vS) * dx(stokes.id)\n",
    "B1S = - qS * div(uS) * dx(stokes.id)\n",
    "B1Dt = - pD * div(vD) * dx(darcy.id)\n",
    "B1D = - qD * div(uD) * dx(darcy.id)\n",
    "\n",
    "B2St = avg(la) * dot(vS(\"+\"), n(\"+\")) * dS(interface.id)\n",
    "B2S = avg(xi) * dot(uS(\"+\"), n(\"+\")) * dS(interface.id)\n",
    "B2Dt = avg(la) * dot(vD(\"-\"), n(\"-\")) * dS(interface.id)\n",
    "B2D = avg(xi) * dot(uD(\"-\"), n(\"-\")) * dS(interface.id)\n",
    "\n",
    "FuS = dot(fS, vS) * dx(stokes.id)\n",
    "FuD = dot(fD, vD) * dx(darcy.id)\n",
    "GqS = - gS * qS * dx(stokes.id)\n",
    "GqD = - gD * qD * dx(darcy.id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lhs = [[AS, None, B1St, None, B2St],\n",
    "       [None, AD, None, B1Dt, B2Dt],\n",
    "       [B1S, None, None, None, None],\n",
    "       [None, B1D, None, None, None],\n",
    "       [B2S, B2D, None, None, None]]\n",
    "rhs = [FuS, FuD, GqS, GqD, None]\n",
    "\n",
    "lhs[0][0] += Constant(mesh, 0.) * inner(uS, vS) * dx\n",
    "lhs[1][1] += Constant(mesh, 0.) * inner(uD, vD) * dx\n",
    "lhs[2][2] = Constant(mesh, 0.) * inner(pS, qS) * dx\n",
    "lhs[3][3] = Constant(mesh, 0.) * inner(pD, qD) * dx\n",
    "lhs[4][4] = Constant(mesh, 0.) * avg(la) * avg(xi) * dS(interface.id)\n",
    "rhs[-1] = Constant(mesh, 0.) * avg(xi) * dS(interface.id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def inflow_eval(x):\n",
    "    values = np.zeros((2, x.shape[1]))\n",
    "    values[1, :] = x[0, :]**2 - 1.0\n",
    "    return values\n",
    "\n",
    "\n",
    "inflow = Function(P2v)\n",
    "inflow.interpolate(inflow_eval)\n",
    "noSlip0 = Function(P2v)\n",
    "noSlip1 = Function(BDM1)\n",
    "\n",
    "bdofs_inlet = locate_dofs_topological(P2v, mesh.topology.dim - 1, boundaries_entities[inlet.id])\n",
    "bdofs_wallS = locate_dofs_topological(P2v, mesh.topology.dim - 1, boundaries_entities[wall_stokes.id])\n",
    "bdofs_wallD = locate_dofs_topological(BDM1, mesh.topology.dim - 1, boundaries_entities[wall_darcy.id])\n",
    "bcUin = DirichletBC(inflow, bdofs_inlet)\n",
    "bcUS = DirichletBC(noSlip0, bdofs_wallS)\n",
    "bcUD = DirichletBC(noSlip1, bdofs_wallD)\n",
    "bcs = [bcUin, bcUS, bcUD]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solve multiphysics system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the block linear system\n",
    "A = assemble_matrix_block(lhs, bcs=bcs, restriction=(restriction, restriction))\n",
    "A.assemble()\n",
    "F = assemble_vector_block(rhs, lhs, bcs=bcs, restriction=restriction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "solution = create_vector_block(rhs, restriction=restriction)\n",
    "ksp = PETSc.KSP()\n",
    "ksp.create(mesh.mpi_comm())\n",
    "ksp.setOperators(A)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F, solution)\n",
    "solution.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the block solution in components\n",
    "uS_h, uD_h, pS_h, pD_h, la_h = [Function(V) for V in (P2v, BDM1, P1, P0, Pt)]\n",
    "with BlockVecSubVectorWrapper(solution, [V.dofmap for V in (P2v, BDM1, P1, P0, Pt)],\n",
    "                              restriction) as solution_wrapper:\n",
    "    for solution_wrapper_local, component in zip(solution_wrapper, (uS_h, uD_h, pS_h, pD_h, la_h)):\n",
    "        with component.vector.localForm() as component_local:\n",
    "            component_local[:] = solution_wrapper_local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert np.isclose(uS_h.vector.norm(PETSc.NormType.NORM_2), 73.46630)\n",
    "assert np.isclose(uD_h.vector.norm(PETSc.NormType.NORM_2), 2.709167)\n",
    "assert np.isclose(pS_h.vector.norm(PETSc.NormType.NORM_2), 174.8691)\n",
    "assert np.isclose(pD_h.vector.norm(PETSc.NormType.NORM_2), 54.34432)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
